import os
from dotenv import load_dotenv

from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent

from my_deep_agents_from_scratch.state import DeepAgentState
from my_deep_agents_from_scratch.research_tools import tavily_search, think_tool, get_today_str
from my_deep_agents_from_scratch.task_tool import _create_task_tool
from my_deep_agents_from_scratch.todo_tools import write_todos, read_todos
from my_deep_agents_from_scratch.file_tools import ls, read_file, write_file
from utils import format_messages, show_prompt

from datetime import datetime

def main():
    load_dotenv(os.path.join("..", ".env"), override=True)
    # model = os.getenv("MODEL")
    # base_url = os.getenv("BASE_URL")
    # api_key = os.getenv("_API_KEY")
    model = os.getenv("VLLM_MODEL")
    base_url = os.getenv("VLLM_BASE_URL")
    api_key = os.getenv("VLLM_API_KEY")

    from my_deep_agents_from_scratch.prompts import (
        FILE_USAGE_INSTRUCTIONS,
        RESEARCHER_INSTRUCTIONS,
        SUBAGENT_USAGE_INSTRUCTIONS,
        TODO_USAGE_INSTRUCTIONS,
    )

    # Create agent using create_react_agent directly
    # from langchain_deepseek.chat_models import ChatDeepSeek
    # model = ChatDeepSeek(model=model, temperature=0, api_key=api_key, base_url=base_url)
    model = ChatOpenAI(model=model, base_url=base_url, api_key=api_key)

    # Limits
    max_concurrent_research_units = 3
    max_researcher_iterations = 3

    # Tools
    sub_agent_tools = [tavily_search, think_tool]
    built_in_tools = [ls, read_file, write_file, write_todos, read_todos, think_tool]

    # Create research sub-agent
    research_sub_agent = {
        "name": "research-agent",
        "description": "Delegate research to the sub-agent researcher. Only give this researcher one topic at a time.",
        "prompt": RESEARCHER_INSTRUCTIONS.format(date=get_today_str()),
        "tools": ["tavily_search", "think_tool"],
    }

    # Create task tool to delegate tasks to sub-agents
    task_tool = _create_task_tool(
        sub_agent_tools, [research_sub_agent], model, DeepAgentState
    )

    delegation_tools = [task_tool]
    all_tools = sub_agent_tools + built_in_tools + delegation_tools  # search available to main agent for trivial cases

    # Build prompt
    SUBAGENT_INSTRUCTIONS = SUBAGENT_USAGE_INSTRUCTIONS.format(
        max_concurrent_research_units=max_concurrent_research_units,
        max_researcher_iterations=max_researcher_iterations,
        date=datetime.now().strftime("%a %b %d, %Y"),
    )
    INSTRUCTIONS = (
            "# TODO MANAGEMENT\n"
            + TODO_USAGE_INSTRUCTIONS
            + "\n\n"
            + "=" * 80
            + "\n\n"
            + "# FILE SYSTEM USAGE\n"
            + FILE_USAGE_INSTRUCTIONS
            + "\n\n"
            + "=" * 80
            + "\n\n"
            + "# SUB-AGENT DELEGATION\n"
            + SUBAGENT_INSTRUCTIONS
    )
    # show_prompt(INSTRUCTIONS)
    # Create agent
    agent = create_react_agent(model, all_tools, prompt=INSTRUCTIONS, state_schema=DeepAgentState)
    result = agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "给我一份模型上下文协议（MCP）的简要概述。",
                }
            ],
        }
    )

    format_messages(result["messages"])


if __name__ == '__main__':
    main()